package ITS;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Signature;
import java.security.SignatureException;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.spec.SecretKeySpec;

public class RSF {

	public static void main(String[] args) {

		if (args.length != 4) {
			System.err.println("NO Valid Arguments!!!");
			System.exit(-1);
		}
		String fileOURPrivateKey = args[0];
		String fileTHEIRPublicKey = args[1];
		String fileToDecrypt = args[2];
		String fileToWrite = args[3];

		PublicKey publicKey = ReadKeys.readPublicKey(fileTHEIRPublicKey);
		PrivateKey privateKey = ReadKeys.readPrivateKey(fileOURPrivateKey);

		byte[] key = null, signature = null/*, file = null*/;

		try {
			File f = new File(fileToDecrypt);
			FileInputStream fis = new FileInputStream(f);
			DataInputStream dis = new DataInputStream(fis);

			int lKey = dis.readInt();
			key = new byte[lKey];
			dis.read(key);

			int lSignature = dis.readInt();
			signature = new byte[lSignature];
			dis.read(signature);

			SecretKeySpec aesKey = decryptAESKey(key, privateKey);
			boolean verifyed = verifySignature(signature, aesKey.getEncoded(), publicKey);
			if (!verifyed){
				System.err.println("Not a valid Signature!\nExiting");
//				System.exit(-1);
			}
			
			decryptFileAndWrite(dis, aesKey, fileToWrite);

//			file = new byte[(dis.available())];
//			dis.read(file);

			dis.close();
		} catch (IOException e) {
			e.printStackTrace();
		}
		
		System.out.println("Done");

	}

	private static boolean verifySignature(byte[] signature, byte[] aesKey,
			PublicKey publicKey) {
		
		boolean ok = false;
		try {
			Signature rsa = Signature.getInstance("SHA1withRSA");
			rsa.initVerify(publicKey);
			rsa.update(aesKey);
			ok = rsa.verify(signature);

		} catch (NoSuchAlgorithmException e) {
			e.printStackTrace();
		} catch (InvalidKeyException e) {
			e.printStackTrace();
		} catch (SignatureException e) {
			e.printStackTrace();
		}
		return ok;
	}

	public static SecretKeySpec decryptAESKey(byte[] key, PrivateKey privateKey) {
		SecretKeySpec skspec = null;
		try {
			Cipher cipher = Cipher.getInstance("RSA");

			cipher.init(Cipher.DECRYPT_MODE, privateKey);
			
			byte[] plain = key;

//			byte[] decData = cipher.update(plain);

			byte[] decRest = cipher.doFinal(plain);

//			byte[] secKey = (/*new String(decData) + */new String(decRest))
//					.getBytes();

			skspec = new SecretKeySpec(decRest, "AES");
		} catch (NoSuchAlgorithmException e) {
			e.printStackTrace();
		} catch (InvalidKeyException e) {
			e.printStackTrace();
		} catch (IllegalBlockSizeException e) {
			e.printStackTrace();
		} catch (BadPaddingException e) {
			e.printStackTrace();
		} catch (NoSuchPaddingException e) {
			e.printStackTrace();
		}
		return skspec;
	}

	public static void decryptFileAndWrite(DataInputStream dis, SecretKeySpec aesKey,
			String outputFile) {
		try {
			Cipher cipher = Cipher.getInstance("AES");

			cipher.init(Cipher.DECRYPT_MODE, aesKey);

			byte[] file = new byte[1024];
			int len = 0;

			DataOutputStream os = new DataOutputStream(new FileOutputStream(
					outputFile));
			
			while ((len = dis.read(file)) > 0) {
			
				///*byte[] decData = */cipher.update(file);
				os.write(cipher.update(file, 0 , len));

			}

			//byte[] decRest = cipher.doFinal();
			os.write(cipher.doFinal());

			//byte[] decrFile = decRest; //(new String(decData) + new String(decRest))
					//.getBytes();

			os.close();
		} catch (IOException e) {
			e.printStackTrace();
		} catch (NoSuchAlgorithmException e) {
			e.printStackTrace();
		} catch (NoSuchPaddingException e) {
			e.printStackTrace();
		} catch (InvalidKeyException e) {
			e.printStackTrace();
		} catch (IllegalBlockSizeException e) {
			e.printStackTrace();
		} catch (BadPaddingException e) {
			e.printStackTrace();
		}
	}

}
